#ifndef AIPOSTPROCESS_H
#define AIPOSTPROCESS_H

#include "inc/rknn_api.h"
#include "aiimageutils.h"
#include <iostream>
#include <vector>
#include <algorithm>
#include <numeric>
#include <opencv2/opencv.hpp>

#if 1
#define OBJ_NUMB_MAX_SIZE 32
typedef struct {
    rknn_context rknn_ctx;
    rknn_input_output_num io_num;
    rknn_tensor_attr* input_attrs;
    rknn_tensor_attr* output_attrs;
    int model_channel;
    int model_width;
    int model_height;
    bool is_quant;
} rknn_app_context_t;

typedef struct {
    image_rect_t box;
    float prop;
    int cls_id;
} object_detect_result;

typedef struct {
    int id;
    int count;
    object_detect_result results[OBJ_NUMB_MAX_SIZE];
} object_detect_result_list;

class AiPostProcess {
public:
    //static const int OBJ_CLASS_NUM = 6;
    int objClassNum;
    int objNumbMaxSize;
    float nmsThresh;
    float boxThresh;

    AiPostProcess(int objNumbMaxSize_ = 64, float nmsThresh_ = 0.35,
                  float boxThresh_ = 0.25,int objClassNum_= 11 );

    int initPostProcess(const std::string &labelNameTxTPath, int objClassNum_);
    void deinitPostProcess();
    char* cocoClsToName(int cls_id);
    int postProcess(rknn_app_context_t *appCtx, rknn_tensor_mem **outputs, letterbox_t *letterBox, float confThreshold, float nmsThreshold, object_detect_result_list *odResults);

    void setnmsThresh(float nmsThresh_){nmsThresh=nmsThresh_;};
    void setboxThresh(float boxThresh_){boxThresh = boxThresh_ ; } ;

private:
    char *labels[OBJ_NUMB_MAX_SIZE];

    static char* readLine(FILE *fp, char *buffer, int *len);
    static int readLines(const char *fileName, char *lines[], int maxLine);
    int loadLabelName(const char *locationFilename, char *label[], int objClassNum);
    static float calculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1, float ymax1);
    static int nms(int validCount, std::vector<float> &outputLocations, std::vector<int> classIds, std::vector<int> &order, int filterId, float threshold);
    static int quickSortIndiceInverse(std::vector<float> &input, int left, int right, std::vector<int> &indices);
    static float sigmoid(float x);
    static float unsigmoid(float y);
    static int32_t clip(float val, float min, float max);
    static int8_t qntF32ToAffine(float f32, int32_t zp, float scale);
    static float deqntAffineToF32(int8_t qnt, int32_t zp, float scale);
    void computeDfl(float* tensor, int dflLen, float* box);
    int processI8(int8_t *boxTensor, int32_t boxZp, float boxScale, int8_t *scoreTensor, int32_t scoreZp, float scoreScale,
                  int8_t *scoreSumTensor, int32_t scoreSumZp, float scoreSumScale, int gridH, int gridW, int stride, int dflLen,
                  std::vector<float> &boxes, std::vector<float> &objProbs, std::vector<int> &classId, float threshold);
    int processFp32(float *boxTensor, float *scoreTensor, float *scoreSumTensor, int gridH, int gridW, int stride, int dflLen,
                    std::vector<float> &boxes, std::vector<float> &objProbs, std::vector<int> &classId, float threshold);
};





#else
#define OBJ_THRESH 0.5
#define NMS_THRESH 0.45
#define IMG_SIZE {640, 480}

void LogDebug(const std::string& message) {
    std::cout << "DEBUG: " << message << std::endl;
}

struct Box {
    float x, y, w, h;
};

struct Attribute {
    int index;
    std::string name;
    int n_dims;
    std::vector<int> dims;
    int n_elems;
    int size;
    int w_stride;
    int size_with_stride;
    std::string fmt;
    std::string type;
    std::string qnt_type;
    int zp;
    float scale;
};

// 过滤框函数
std::tuple<std::vector<Box>, std::vector<int>, std::vector<float>> filter_boxes(
    const std::vector<Box>& boxes,
    const std::vector<float>& box_confidences,
    const std::vector<std::vector<float>>& box_class_probs) {

    std::vector<float> class_max_score;
    std::vector<int> classes;
    std::vector<float> scores;

    for (size_t i = 0; i < box_class_probs.size(); ++i) {
        auto max_it = std::max_element(box_class_probs[i].begin(), box_class_probs[i].end());
        class_max_score.push_back(*max_it);
        classes.push_back(std::distance(box_class_probs[i].begin(), max_it));
    }

    std::vector<int> _class_pos;
    for (size_t i = 0; i < class_max_score.size(); ++i) {
        if (class_max_score[i] * box_confidences[i] >= OBJ_THRESH) {
            _class_pos.push_back(i);
            scores.push_back(class_max_score[i] * box_confidences[i]);
        }
    }

    std::vector<Box> filtered_boxes;
    std::vector<int> filtered_classes;

    for (auto idx : _class_pos) {
        filtered_boxes.push_back(boxes[idx]);
        filtered_classes.push_back(classes[idx]);
    }

    return {filtered_boxes, filtered_classes, scores};
}

// 非极大值抑制函数
std::vector<int> nms_boxes(const std::vector<Box>& boxes, const std::vector<float>& scores) {
    std::vector<int> keep;
    std::vector<int> order(scores.size());
    std::iota(order.begin(), order.end(), 0);
    std::sort(order.begin(), order.end(), [&scores](int a, int b) { return scores[a] > scores[b]; });

    while (!order.empty()) {
        int i = order.front();
        keep.push_back(i);
        std::vector<int> remaining;

        for (size_t j = 1; j < order.size(); ++j) {
            int k = order[j];
            float xx1 = std::max(boxes[i].x, boxes[k].x);
            float yy1 = std::max(boxes[i].y, boxes[k].y);
            float xx2 = std::min(boxes[i].x + boxes[i].w, boxes[k].x + boxes[k].w);
            float yy2 = std::min(boxes[i].y + boxes[i].h, boxes[k].y + boxes[k].h);

            float w = std::max(0.0f, xx2 - xx1 + 0.00001f);
            float h = std::max(0.0f, yy2 - yy1 + 0.00001f);
            float inter = w * h;
            float ovr = inter / (boxes[i].w * boxes[i].h + boxes[k].w * boxes[k].h - inter);

            if (ovr <= NMS_THRESH) {
                remaining.push_back(k);
            }
        }
        order = remaining;
    }
    return keep;
}

// DFL 函数
std::vector<cv::Mat> dfl(const std::vector<cv::Mat>& positions) {
    // 实现 DFL 逻辑（类似于 Python 代码中使用的 torch 操作）
    std::vector<cv::Mat> results;
    for (const auto& position : positions) {
        // 添加 DFL 实现细节...
    }
    return results;
}

// 处理框函数
std::vector<cv::Mat> box_process(const cv::Mat& position) {
    int grid_h = position.size[2];
    int grid_w = position.size[3];

    cv::Mat col, row;
    cv::meshgrid(cv::range(0, grid_w), cv::range(0, grid_h), col, row);
    cv::Mat grid = cv::Mat(2, 1, CV_32F);
    cv::hconcat(col.reshape(1, 1), row.reshape(1, 1), grid);

    cv::Mat stride = (cv::Mat_<float>(2, 1) << IMG_SIZE[1] / grid_h, IMG_SIZE[0] / grid_w);
    auto positions_dfl = dfl(position);

    // Implement box processing logic...
    // ...

    std::vector<cv::Mat> xyxy;
    return xyxy;
}

// 后处理函数
std::tuple<std::vector<Box>, std::vector<int>, std::vector<float>> post_process(const std::vector<cv::Mat>& input_data) {
    std::vector<std::vector<Box>> boxes;
    std::vector<std::vector<float>> scores, classes_conf;

    int defualt_branch = 3;
    int pair_per_branch = input_data.size() / defualt_branch;

    for (int i = 0; i < defualt_branch; ++i) {
        boxes.push_back(box_process(input_data[pair_per_branch * i]));
        classes_conf.push_back(input_data[pair_per_branch * i + 1]);
        scores.push_back(cv::Mat::ones(input_data[pair_per_branch * i + 1].size[1], CV_32F));
    }

    // Implement flattening and concatenation logic...

    std::tie(boxes, classes, scores) = filter_boxes(boxes, scores, classes_conf);

    // NMS 逻辑
    std::vector<Box> nboxes;
    std::vector<int> nclasses;
    std::vector<float> nscores;

    for (const auto& c : unique(classes)) {
        auto inds = std::find(classes.begin(), classes.end(), c);
        auto b = boxes[inds];
        auto s = scores[inds];
        auto keep = nms_boxes(b, s);

        if (!keep.empty()) {
            nboxes.insert(nboxes.end(), b[keep.begin(), keep.end()]);
            nclasses.insert(nclasses.end(), c[keep.begin(), keep.end()]);
            nscores.insert(nscores.end(), s[keep.begin(), keep.end()]);
        }
    }

    if (nclasses.empty() && nscores.empty()) {
        return std::make_tuple(std::vector<Box>(), std::vector<int>(), std::vector<float>());
    }

    boxes = nboxes;
    classes = nclasses;
    scores = nscores;

    return std::make_tuple(boxes, classes, scores);
}

int main() {
    LogDebug("11111");
    // 添加你的测试代码来验证功能...
    return 0;
}
#endif
#endif
